import keras
from keras.applications.mobilenet import MobileNet, BASE_WEIGHT_PATH, get_file, relu6, DepthwiseConv2D
from keras_retinanet.models import retinanet
from new_model.segmentation_branch import segmentation_branch
from new_model.config import TrainingConfig, ModelConfig


def model_generator(backbone='mobilenet224_1.0', inputs=None, **kwargs):
    alpha = float(backbone.split('_')[1])
    # choose default input
    if inputs is None:
        inputs = keras.layers.Input(ModelConfig.input_shape)
    # load backbone
    mobilenet = MobileNet(input_tensor=inputs, alpha=alpha, include_top=False, pooling=None, weights='imagenet')
    layer_names = ['conv_pw_5_relu', 'conv_pw_11_relu', 'conv_pw_13_relu'] # stride size = [8, 16, 32]
    layer_outputs = [mobilenet.get_layer(name).output for name in layer_names]
    mobilenet = keras.models.Model(inputs=inputs, outputs=layer_outputs, name=mobilenet.name)
    # freeze backbone
    for layer in mobilenet.layers:
        layer.trainable = TrainingConfig.freeze_backbone
    # build detection branch
    detection_model = retinanet.retinanet(inputs=inputs, num_classes=ModelConfig.detection_num_classes,
                                          backbone_layers=mobilenet.outputs, name='detection_model', **kwargs)
    # build segmentation branch
    segmentation_model = keras.models.Model(inputs=inputs, outputs=segmentation_branch(mobilenet.outputs[0], num_classes=ModelConfig.segmentation_num_classes),
                                            name='segmentation_model')
    # united model
    unit_model = keras.models.Model(inputs=inputs, outputs=detection_model.outputs[:2] + segmentation_model.outputs)
    return detection_model, segmentation_model, unit_model


if __name__ == '__main__':
    d, s, u = model_generator()
    # visualize network with keras api
    from keras.utils import plot_model
    import os
    os.environ["PATH"] += os.pathsep + 'C:/Program Files (x86)/Graphviz2.38/bin/'
    plot_model(d, show_shapes=True, to_file='detection.png')
    plot_model(s, show_shapes=True, to_file='segmentation.png')
    plot_model(u, show_shapes=True, to_file='unite_model.png')